package proxy

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"sync"
)

// http200 http 200 response
var http200 = []byte("HTTP/1.1 200 Connection Established\r\n\r\n")

func New() http.Handler {
	return &Proxy{
		tr: &http.Transport{DisableKeepAlives: true},
		bp: &sync.Pool{New: func() interface{} { return make([]byte, 1<<12) }},
	}
}

type Proxy struct {
	tr *http.Transport
	bp *sync.Pool
	vb bool
}

func (proxy *Proxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	defer func() {
		if err := recover(); err != nil {
			rw.WriteHeader(http.StatusInternalServerError)
			log.Printf("panic: %v\n", err)
			_, _ = fmt.Fprint(rw, err)
		}
	}()

	if req.Method == "CONNECT" {
		proxy.handleHttps(rw, req)
	} else {
		proxy.handleHttp(rw, req)
	}
}

func (proxy *Proxy) handleHttp(rw http.ResponseWriter, req *http.Request) {
	if proxy.vb {
		log.Printf("sending request %v %v \n", req.Method, req.URL.Host)
	}
	removeHopByHop(req)

	resp, err := proxy.tr.RoundTrip(req)
	if err != nil {
		log.Printf("%v", err)
		http.Error(rw, err.Error(), 500)
		return
	}
	defer tryClose(resp.Body)

	headerClear(rw.Header())
	headerCopy(rw.Header(), resp.Header)

	rw.WriteHeader(resp.StatusCode) //写入响应状态

	nr, err := proxy.ioCopy(rw, resp.Body)
	if err != nil && err != io.EOF {
		log.Printf("got an error when copy remote response to client. %v\n", err)
		return
	}
	if proxy.vb {
		log.Printf("copied %v bytes from %v.\n", nr, req.URL.Host)
	}
}

func (proxy *Proxy) handleHttps(rw http.ResponseWriter, req *http.Request) {
	if proxy.vb {
		log.Printf("tried to connect to %v", req.URL.Host)
	}

	hj, _ := rw.(http.Hijacker)
	client, _, err := hj.Hijack() //获取客户端与代理服务器的tcp连接
	if err != nil {
		log.Printf("failed to get Tcp connection of %s \n", req.RequestURI)
		http.Error(rw, "Failed", http.StatusBadRequest)
		return
	}

	remote, err := net.Dial("tcp", req.URL.Host) //建立服务端和代理服务器的tcp连接
	if err != nil {
		log.Printf("failed to connect %v\n", req.RequestURI)
		tryClose(client)
		return
	}

	_, _ = client.Write(http200)

	go func() {
		if _, err := proxy.ioCopy(remote, client); err != nil {
			log.Println(err)
		}
	}()

	go func() {
		if _, err := proxy.ioCopy(client, remote); err != nil {
			log.Println(err)
		}
	}()
}

func (proxy *Proxy) ioCopy(dst io.Writer, src io.ReadCloser) (nr int, err error) {
	defer tryClose(src)

	var buf = proxy.bp.Get().([]byte)
	defer proxy.bp.Put(buf)

	var re, we error
	var n int
	for {
		if n, re = src.Read(buf); n > 0 {
			_, we = dst.Write(buf[:n])
			tryFlush(dst)
			nr += n
		}

		if err = re; we != nil {
			err = we
		}
		if err != nil {
			if err == io.EOF {
				err = nil
			}
			return
		}
	}
}

func tryClose(closer io.Closer) {
	if closer != nil {
		if err := closer.Close(); err != nil {
			log.Println(err)
		}
	}
}

func tryFlush(w io.Writer) {
	if flusher, ok := w.(http.Flusher); ok {
		flusher.Flush()
	}
}
